Conversation
Capture WAIT/reset and H2D during actual FULL graph capture, release pending fills before synchronization, and fall back on unsupported stream memops. Include an in-tree CUDA driver binding and lifecycle/graph tests. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
|
Review 1 (claude_vllm_freetoken), head 44e4012 vs base a7a3d7a (a97 + vllm-project#54129), production diff read: ple_layer.py (custom op + DeferredRows init/prepare/dummy), model_state.py (step eligibility, complete/abort, prepare_inputs branch), model_runner.py (eligibility set before prepare_inputs, complete after run_fullgraph, abort on exceptions), ple_wait.py (DeferredRows), csrc/ple_memops.cpp + CMake/setup, tests skimmed. No blocking findings. The capture condition covers FULL capture running as mode NONE (the defect found on the lab branch), scoped to a single real token; eligibility is set per step before prepare_inputs and cleared in complete(); dummy runs pre-signal via prepare_dummy so captured WAITs never strand; failures abort and poison permanently (stated); the memops extension builds only for CUDA and the CUDA test skips when the stream memops are unavailable. One question before approval: availability is probed per DeferredRows at initialize_mmap_staging, with a fallback warning per module. If one module's probe fails and another's succeeds, set_deferred_ple_step() is False for every step (all-or-none), but the modules that did get DeferredRows still capture the WAIT/H2D in the single-token FULL graph and the synchronous prepare_mmap_rows path never signals their flag, so the first replay would wait forever. Is a mixed state impossible by construction (same device, same probe)? If so a short comment or an all-or-none assert at the end of initialize would close this; otherwise the sync path should signal when deferred_rows exists. Nit: qwen4_exp_ple_deferred_rows dereferences layer.ple_embedding.deferred_rows without a None check inside the capture branch; forward() only calls the op when it is not None, so this is fine as long as that invariant holds. Not executed by me: CPU tests (author reports 12 + 5 subtests pass, pre-commit/mypy/clang pass); CUDA build, smoke and checkpoint on the new base pending astra. |
Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
|
Review 1 (claude_vllm_freetoken): head 650be6c resolves the question from my first pass — if any mmap layer's stream-memops probe fails, initialize_mmap_staging now drops every layer's DeferredRows (all-or-none) with a warning, before any prepare or capture can run, so no layer can capture a WAIT/H2D that the synchronous path would never signal. Approve from review 1 (code/tests read; CPU tests per author 13 passed / 5 subtests / 1 CUDA skip; CUDA build, smoke and checkpoint on the PLE base pending astra). Merge is the author's call after review 2. |
Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
1 participant
Enable opt-in deferred mmap PLE row loading for single-token FULL CUDA graphs. The graph waits for a host gather, resets its flag, and copies the ready row into stable staging. Capture explicitly handles the NONE runtime mode used inside FULL graph capture; eager and PIECEWISE execution retain synchronous preparation.
Dependency and scope:
Implementation: in-tree optional CUDA driver memops extension, pinned-buffer lifetime and capability probe, synchronous fallback for unsupported memops, completion and exception-release hooks in the consuming runner, raw-byte lifecycle tests, and an actual CUDA graph WAIT/reset/H2D test. Other driver errors remain errors. Abort poisons the helper rather than silently reusing failed data.
Validation at head
4f859de9d0f55760b50358aee4834e6966e13bc8:_ple_memops, component install, and candidate import passed with the upstream a97 wheel / torch 2.13.0+cu130 environment. CUDA lifecycle/graph tests: 2 passed, no skips, 3.65 s. Evidence:results/upstream-ple-exact-4f859de9d0f5-20260909/(configure/build/install/import logs, pytest XML, binary provenance).results/upstream-ple-checkpoint-4f859de9d0f5-deferred{0,1}-20260909/.Integration status (separate from this PR's standalone smoke):
7dedc6d8d9b178b60f6a5b32f03d677145982441includes the standalone expert pool plus this PLE dependency stack. On 96 GiB Blackwell ballooned to 48 GiB available, Marlin/no-speculation/chunk512/FULL_DECODE_ONLY/deferred ON, speed context 4096: fresh three-run B rates 63.188228 / 62.708682 / 63.651909 tok/s; median 63.188228, passing the preregistered 62.985 threshold. These are integrated results, not this PR's isolated speed contribution or Ada 48 GB measurements. Evidence:results/goal66-minimal-repeat3-20260909/MID-FINAL-REVIEW.json.results/goal66-minimal-quality-20260909/MID-FINAL-QUALITY-REVIEW.mdand per-task reviews. Final state exit 0/not OOM-killed; 29 startup + 243 quality capacity samples had no guard error. Startup allocator OOM retry and shutdown warnings are retained; samples do not establish continuous peak usage. Earlier assembly results were not inherited.Delta: 10 files, +1199/-40 against
lab/upstream-ple-mmap-base. AI-assisted implementation. Author: mid; independent review 1: Claude; independent review 2 and GPU execution: Astra. Both reviewers' reported standalone findings were addressed; final minimal integration measurement acceptance is recorded with the limitations above. This PR remains Draft; no ready/merge action is included. No head changes in this documentation update. Fork-only preparation; no upstream submission or related-issue posting.